{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Analyzing the Parker Solar Probe flybys\n\n\n 1. Modulus of the exit velocity, some features of Orbit #2\n ----------------------------------------------------------\n\nFirst, using the data available in the reports, we try to compute some of the properties of orbit #2.\nThis is not enough to completely define the trajectory, but will give us information later on in the process.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from astropy import units as u\n\nT_ref = 150 * u.day\nprint(T_ref)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from poliastro.bodies import Earth, Sun, Venus\n\nk = Sun.k\nprint(k)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\\begin{align}T = 2 \\pi \\sqrt{\\frac{a^3}{\\mu}} \\Rightarrow a = \\sqrt[3]{\\frac{\\mu T^2}{4 \\pi^2}}\\end{align}\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "a_ref = np.cbrt(k * T_ref**2 / (4 * np.pi**2)).to(u.km)\nprint(a_ref.to(u.au))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\\begin{align}\\varepsilon = -\\frac{\\mu}{r} + \\frac{v^2}{2} = -\\frac{\\mu}{2a} \\Rightarrow v = +\\sqrt{\\frac{2\\mu}{r} - \\frac{\\mu}{a}}\\end{align}\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "energy_ref = (-k / (2 * a_ref)).to(u.J / u.kg)\nprint(energy_ref)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from poliastro.twobody import Orbit\nfrom poliastro.util import norm\nfrom astropy.time import Time\n\nflyby_1_time = Time(\"2018-09-28\", scale=\"tdb\")\nprint(flyby_1_time)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "r_mag_ref = norm(Orbit.from_body_ephem(Venus, epoch=flyby_1_time).r)\nprint(r_mag_ref.to(u.au))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "v_mag_ref = np.sqrt(2 * k / r_mag_ref - k / a_ref)\nprint(v_mag_ref.to(u.km / u.s))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "2. Lambert arc between #0 and #1\n--------------------------------\n\nTo compute the arrival velocity to Venus at flyby #1, we have the necessary data to solve the boundary value problem.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "d_launch = Time(\"2018-08-11\", scale=\"tdb\")\nprint(d_launch)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ss0 = Orbit.from_body_ephem(Earth, d_launch)\nss1 = Orbit.from_body_ephem(Venus, epoch=flyby_1_time)\n\ntof = flyby_1_time - d_launch\n\nfrom poliastro import iod\n\n(v0, v1_pre), = iod.lambert(Sun.k, ss0.r, ss1.r, tof.to(u.s))\nprint(v0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(v1_pre)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(norm(v1_pre))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "3. Flyby #1 around Venus\n------------------------\n\nWe compute a flyby using poliastro with the default value of the entry angle,\njust to discover that the results do not match what we expected.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from poliastro.threebody.flybys import compute_flyby\n\nV = Orbit.from_body_ephem(Venus, epoch=flyby_1_time).v\nprint(V)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "h = 2548 * u.km\n\nd_flyby_1 = Venus.R + h\nprint(d_flyby_1.to(u.km))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "V_2_v_, delta_ = compute_flyby(v1_pre, V, Venus.k, d_flyby_1)\n\nprint(norm(V_2_v_))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "4. Optimization\n---------------\n\nNow we will try to find the value of $\\theta$ that satisfies our requirements.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def func(theta):\n    V_2_v, _ = compute_flyby(v1_pre, V, Venus.k, d_flyby_1, theta * u.rad)\n    ss_1 = Orbit.from_vectors(Sun, ss1.r, V_2_v, epoch=flyby_1_time)\n    return (ss_1.period - T_ref).to(u.day).value"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "There are two solutions:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n# %matplotlib inline\n\ntheta_range = np.linspace(0, 2 * np.pi)\nplt.figure()\nplt.plot(theta_range, [func(theta) for theta in theta_range])\nplt.axhline(0, color='k', linestyle=\"dashed\")\nplt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(func(0))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(func(1))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from scipy.optimize import brentq\n\ntheta_opt_a = brentq(func, 0, 1) * u.rad\nprint(theta_opt_a.to(u.deg))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "theta_opt_b = brentq(func, 4, 5) * u.rad\nprint(theta_opt_b.to(u.deg))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "V_2_v_a, delta_a = compute_flyby(v1_pre, V, Venus.k, d_flyby_1, theta_opt_a)\nV_2_v_b, delta_b = compute_flyby(v1_pre, V, Venus.k, d_flyby_1, theta_opt_b)\n\nprint(norm(V_2_v_a))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(norm(V_2_v_b))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "5. Exit Orbit\n---------------\n\nAnd finally, we compute orbit #2 and check that the period is the expected one.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ss01 = Orbit.from_vectors(Sun, ss1.r, v1_pre, epoch=flyby_1_time)\nprint(ss01)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The two solutions have different inclinations, so we still have to find out which is the good one.\nWe can do this by computing the inclination over the ecliptic - however, as the original data was\nin the International Celestial Reference Frame (ICRF), whose fundamental plane is parallel to the\nEarth equator of a reference epoch, we have change the plane to the Earth ecliptic,\nwhich is what the original reports use.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ss_1_a = Orbit.from_vectors(Sun, ss1.r, V_2_v_a, epoch=flyby_1_time)\nprint(ss_1_a)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ss_1_b = Orbit.from_vectors(Sun, ss1.r, V_2_v_b, epoch=flyby_1_time)\nprint(ss_1_b)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's define a function to do that quickly for us, using the get_frame function from poliastro.frames:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from astropy.coordinates import CartesianRepresentation\nfrom poliastro.frames import Planes, get_frame\n\ndef change_plane(ss_orig, plane):\n    \"\"\"Changes the plane of the Orbit.\n\n    \"\"\"\n    ss_orig_rv = ss_orig.frame.realize_frame(\n        ss_orig.represent_as(CartesianRepresentation)\n    )\n\n    dest_frame = get_frame(ss_orig.attractor, plane, obstime=ss_orig.epoch)\n\n    ss_dest_rv = ss_orig_rv.transform_to(dest_frame)\n    ss_dest_rv.representation_type = CartesianRepresentation\n\n    ss_dest = Orbit.from_vectors(\n        ss_orig.attractor,\n        r=ss_dest_rv.data.xyz,\n        v=ss_dest_rv.data.differentials['s'].d_xyz,\n        epoch=ss_orig.epoch,\n        plane=plane,\n    )\n    return ss_dest\n\nprint(change_plane(ss_1_a, Planes.EARTH_ECLIPTIC))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(change_plane(ss_1_b, Planes.EARTH_ECLIPTIC))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Therefore, the correct option is the first one.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(ss_1_a.period.to(u.day))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(ss_1_a.a)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "And, finally, we plot the solution:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from poliastro.plotting import OrbitPlotter\nplt.figure()\nframe = OrbitPlotter()\n\nframe.plot(ss0, label=Earth)\nframe.plot(ss1, label=Venus)\nframe.plot(ss01, label=\"#0 to #1\")\nframe.plot(ss_1_a, label=\"#1 to #2\");"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}